# from playwright.sync_api import sync_playwright
from datetime import datetime, timedelta
import akshare as ak
import util

if __name__ == "__main__":
    conn, cursor = util.getDb()
    cursor.execute("SELECT * FROM stock_info ")
    dbrows = cursor.fetchall()
    stock_zh_a_spot_em_df = ak.stock_zh_a_spot_em()
    # 只保留所需的列
    selected_columns_df = stock_zh_a_spot_em_df[
        ["代码", "名称", "总市值", "流通市值", "市盈率-动态", "市净率"]
    ]
    # 重命名 DataFrame 列与 SQLite 数据库中的列名一致
    selected_columns_df.columns = [
        "symbol",
        "name",
        "market_cap",
        "circulating_cap",
        "pe_ratio_dynamic",
        "pb_ratio",
    ]

    # 检查并插入数据
    for _, row in selected_columns_df.iterrows():
        symbol_needs_update = any(
            row["symbol"] == dbrow[0] and dbrow[6] is None for dbrow in dbrows
        )
        if not row["symbol"] in {row[0] for row in dbrows}:
            stock_info_df = ak.stock_individual_info_em(symbol=row.symbol)
            name = stock_info_df.loc[
                stock_info_df["item"] == "股票简称", "value"
            ].values[0]
            industry = stock_info_df.loc[
                stock_info_df["item"] == "行业", "value"
            ].values[0]
            listing_date_str = stock_info_df.loc[
                stock_info_df["item"] == "上市时间", "value"
            ].values[0]
            total_shares = stock_info_df.loc[
                stock_info_df["item"] == "总股本", "value"
            ].values[0]
            if listing_date_str and not str(listing_date_str) == "-":
                listing_date = datetime.strptime(str(listing_date_str), "%Y%m%d").date()
                cursor.execute(
                    """
                        INSERT INTO stock_info (
                            symbol, name, market_cap, circulating_cap, pe_ratio_dynamic, 
                            pb_ratio, industry, listing_date, total_shares
                        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                        """,
                    (
                        str(row["symbol"]).zfill(6),
                        name,
                        row["market_cap"],
                        row["circulating_cap"],
                        row["pe_ratio_dynamic"],
                        row["pb_ratio"],
                        industry,
                        listing_date,
                        total_shares,
                    ),
                )
        elif symbol_needs_update:
            stock_info_df = ak.stock_individual_info_em(symbol=symbol)
            name = stock_info_df.loc[
                stock_info_df["item"] == "股票简称", "value"
            ].values[0]
            industry = stock_info_df.loc[
                stock_info_df["item"] == "行业", "value"
            ].values[0]
            listing_date_str = stock_info_df.loc[
                stock_info_df["item"] == "上市时间", "value"
            ].values[0]
            total_shares_str = stock_info_df.loc[
                stock_info_df["item"] == "总股本", "value"
            ].values[0]

            if listing_date_str and not str(listing_date_str) == "-":
                cursor.execute(
                    """
                UPDATE stock
                SET name =?,industry = ?, listing_date = ?, total_shares = ?
                WHERE symbol = ?
            """,
                    (name, industry, listing_date, total_shares, symbol),
                )
    # 提交更改并关闭数据库连接
    conn.commit()
    conn.close()
